#任务收集者
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

from sshcode.plan_execute.model_utils import getLLM

_system_prompt_template = '''
你是一个优秀的助手，根据参考结果，准确的回答用户的问题
不要无关信息
'''

_human_prompt_template = '''
参考结果：
{infos}
查询任务：
{query}
'''

class Gatherer:
    def __init__(self,llm):
        _prompt = ChatPromptTemplate.from_messages([
            ("system",_system_prompt_template),
            ("human",_human_prompt_template)
        ])

        _parser = StrOutputParser()
        self.chain = _prompt | llm | _parser

    def __call__(self, state):
        return self.chain.invoke(state)

if __name__ == '__main__':
    llm = getLLM()
    _gatherer = Gatherer(llm=llm)
    _rt = _gatherer({"query": "2024年法国巴黎奥运会女子10米跳水冠军的家乡在哪里",
                     "infos": ["全红禅是2024年法国巴黎奥运会女子10米跳水的冠军", "她住在四川"], })
    print(_rt)